import os
from typing import Any, Callable, List, Optional, Tuple
from pypdf import PdfReader, PdfWriter
import re
import io
from reportlab.lib.pagesizes import letter
from reportlab.pdfbase.pdfmetrics import stringWidth
from reportlab.pdfgen import canvas
from reportlab.pdfbase.pdfmetrics import stringWidth
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.pdfbase import pdfmetrics
from pypdf.annotations import Link
from dataclasses import dataclass

FONT = "Helvetica"


@dataclass
class PDFMergeConfig:
    # 合并哪些文件下的PDF
    folders: List[str]
    output_pdf: str

    # 排序函数，分别对应folders里面的第1，2，3个文件夹
    # 输入为文件名（不含扩展名），返回文件的序号。
    # 比如文件夹下有"a.pdf", "c.pdf", "b.pdf"这三个文件，想以a->c->b的顺序合并起来
    # 那么就可以这样写排序函数： lambda name: {"a": 1, "c": 2, "b": 3}[name]
    # 返回的序号只要递增即可，具体大小无所谓
    sort_func: List[Optional[Callable[[str], Any]]]

    # PDF书签和目录项的格式化函数。
    # 输入：字符串，等于文件的名字（不含扩展名）；
    # 输出：字符串，作为PDF的目录和书签
    label_formatter: List[Optional[Callable[[str], str]]]

    # 是否要插入空白页
    # 如果要双面打印，并且确保打印后的每一份文档都能够独立分开，那么就需要将此项设置为True
    insert_blank_page: bool = True

    # 生成目录的字体文件
    catalog_font_file: str = ""


@dataclass
class PDFFile:
    name: str
    directory: str
    label: str


def split_title_text(title: str, available_width, split_ratio):

    splitted: List[str] = []
    head_ptr = 0
    tail_ptr = 1
    # 寻找合适的分割点
    while True:
        split_title = title[head_ptr:tail_ptr]
        if stringWidth(split_title + "-", FONT, 12) > available_width * split_ratio:
            splitted.append(split_title + "-")
            head_ptr = tail_ptr
        tail_ptr += 1
        if not tail_ptr < len(title):
            splitted.append(title[head_ptr:])
            break
    if len(splitted) > 1:
        splitted[-1] += "  "
    return splitted


# 定义函数，添加目录页到PDF
def add_catalog_page(bookmarks):
    packet = io.BytesIO()  # 创建内存中的字节包，用于临时存储PDF数据
    c = canvas.Canvas(packet, pagesize=letter)  # 创建一个画布，设置其页面大小为letter
    width, height = letter  # 获取页面宽度和高度
    top_margin = 60  # 设置顶部边距
    bottom_margin = 60  # 设置底部边距
    y_position = height - top_margin  # 计算初始y坐标位置，考虑顶部边距
    c.setFont(FONT, 16)  # 设置目录标题的字体和大小
    c.drawString(280, y_position, "目录")  # 绘制目录标题
    y_position -= 30  # 更新y坐标位置以为目录项留出空间

    c.setFont(FONT, 12)  # 设置目录项的字体和大小
    left_margin = 72  # 左边距
    right_margin = width - 72  # 右边距
    dot_space = 5  # 点之间的间隔
    different_title_spacing = 25  # 不同标题之间的间隔
    same_title_line_spacing = 15  # 同一个标题换行的间隔
    split_ratio = 0.97  # 分割点的位置比例

    title_number = 1  # 初始化标题编号
    pagecount = 0  # 当前在生成目录的页的索引，从0开始
    links = []
    for title, page_number in bookmarks:
        box = [left_margin, 0, right_margin, 0]
        links.append((box, page_number, pagecount))
        # title = title.encode("utf8").decode("unicode_escape")
        title = re.sub(r"\[.*?\]", "", title)  # 删除标题中的方括号及其内容
        title = title.replace("_", " ")  # 将标题中的短横线替换成空格

        # 使用zfill确保编号位数为2，例如"01", "02", ..., "10", "11", ...
        formatted_number = str(title_number).zfill(2)
        title = f"{formatted_number}. {title}"  # 给标题添加编号

        title_number_width = (
            stringWidth(f"{title_number}. ", FONT, 12) + 1
        )  # 计算编号宽度
        title_number += 1  # 更新标题编号

        available_width = right_margin - left_margin - dot_space * 2  # 计算可用宽度
        title_width = stringWidth(title, FONT, 12)
        page_number_str = str(page_number + 1)
        page_number_width = stringWidth(page_number_str, FONT, 12)
        box[1] = y_position + different_title_spacing - 5

        split_line_times = 0
        # last_line_width
        # 判断标题是否需要分割
        if title_width > available_width * split_ratio:
            # split_title = title
            # splitted: List[str] = []
            # print("split_title", )
            title_splitted = split_title_text(title, available_width, split_ratio)
            split_line_times = len(title_splitted) - 1
            # 寻找合适的分割点
            # while True:
            #     if (
            #         stringWidth(split_title + "-", FONT, 12)
            #         > available_width * split_ratio
            #     ):
            #         splitted = split_title.split(" ")
            #     split_title = split_title[:-1]
            # split_title += "-"
            # split_line_times
            last_line_content = title_splitted[-1]
            title_line = title_splitted.pop(0)
            c.drawString(left_margin, y_position, title_line)
            for title_line in title_splitted:
                y_position -= same_title_line_spacing
                c.drawString(left_margin + title_number_width, y_position, title_line)

            # box[1] = y_position
            # 为分割后的标题调整y坐标

            # title = title[len(split_title) - 1 :]
            # # 特别注意：这里不再添加编号宽度，因为编号已经在第一行绘制
            # c.drawString(
            #     left_margin + title_number_width, y_position, title
            # )  # 绘制续行部分标题
        else:
            last_line_content = title
            # 如果不需要分割标题，直接绘制
            c.drawString(left_margin, y_position, title)

        # c.drawString(left_margin, y_position, title)  # 绘制标题
        c.drawRightString(right_margin, y_position, page_number_str)  # 绘制页码

        # 绘制点线连接标题和页码
        dot_line_start = left_margin + stringWidth(last_line_content, FONT, 12) + 10
        dot_line_end = right_margin - page_number_width - 10
        current_position = dot_line_start
        print(split_line_times, dot_line_start, dot_line_end)
        # if split_line_times > 0:
        #     current_position += 10
        while current_position < dot_line_end:
            c.drawString(current_position, y_position, ".")
            current_position += dot_space

        box[3] = y_position - 5
        y_position -= different_title_spacing  # 更新y坐标以为下一个标题留出空间
        if y_position < bottom_margin:  # 检查是否需要翻页
            c.showPage()
            pagecount += 1
            y_position = height - top_margin  # 重置y坐标，考虑顶部边距
            c.setFont(FONT, 12)  # 确保新页面使用正确的字体和大小

    c.save()  # 保存PDF到内存中的字节包
    packet.seek(0)  # 将字节包的指针重置到开始位置
    return PdfReader(packet), links  # 创建并返回一个PDF阅读器对象，包含内存中的PDF数据


def get_files_from_folders(config: PDFMergeConfig) -> List[PDFFile]:
    files = []
    for i, folder in enumerate(config.folders):
        files_in_folder = os.listdir(folder)
        files_in_folder.sort(key=config.sort_func[i])
        label_formatter = config.label_formatter[i]
        if label_formatter is None:
            label_formatter = lambda x: x
        for file in files_in_folder:
            files.append(
                PDFFile(file, folder, label_formatter(os.path.splitext(file)[0]))
            )
    return files


def merge_pdfs(
    config: PDFMergeConfig,
):
    global FONT
    if config.catalog_font_file != "":
        font_name, _ = os.path.splitext(os.path.basename(config.catalog_font_file))
        pdfmetrics.registerFont(TTFont(font_name, config.catalog_font_file))
        FONT = font_name

    pdf_writer = PdfWriter()
    files = get_files_from_folders(config)
    bookmarks = []
    label_formatter = config.label_formatter
    if label_formatter is None:
        label_formatter = lambda x: x
    for file in files:
        filename, folder_path = file.name, file.directory
        if filename.endswith(".pdf"):
            pdf_path = os.path.join(folder_path, filename)
            pdf_reader = PdfReader(pdf_path)
            # print(len(pdf_reader.pages), filename)

            bookmarks.append((file.label, len(pdf_writer.pages)))

            for page in pdf_reader.pages:
                pdf_writer.add_page(page)
            if config.insert_blank_page and len(pdf_reader.pages) % 2 == 1:
                blank_page = PdfReader(
                    os.path.join(os.path.dirname(__file__), "blank.pdf")
                ).pages[0]
                pdf_writer.add_page(blank_page)

    catalog_pdf, links = add_catalog_page(bookmarks)

    for page in reversed(catalog_pdf.pages):
        pdf_writer.insert_page(page)
    pdf_writer.add_outline_item("目录", 0)
    # 获取目录长度之后再添加书签
    for item_text, page in bookmarks:
        pdf_writer.add_outline_item(item_text, page + len(catalog_pdf.pages))

    # Create the annotation and add it
    # Add the line
    for box, page_number, catalog_page in links:
        # print(box, page_number, catalog_page)
        annotation = Link(
            rect=box,
            target_page_index=page_number + len(catalog_pdf.pages),
            # border=[0, 0, 1],
        )
        pdf_writer.add_annotation(page_number=catalog_page, annotation=annotation)

    with open(config.output_pdf, "wb") as output_file:
        pdf_writer.write(output_file)
    print(f"文档合并完成，位于 {config.output_pdf}")
